#include "Dofs.h"
#include "Mesh.h"
#include "Point.h"
#include "TemplateElement.h"
#include <Eigen/Sparse>
#include <Eigen/IterativeLinearSolvers>
#include <Eigen/Dense>
#include <ctime>


/**
 * @brief 整合有限元的流程
 */
class FEM
{
public:
    /**
     * @brief 默认构造函数
     */
    FEM() = default;
    /**
     * @brief 默认析构函数
     */	
    ~FEM() = default;
    void make_grid(const vector<int>&);
    void run(vector<int> _seg);
private:
    /**
     * @brief 生成网格
     */
    
    /**
     * @brief 分配自由度，设置矩阵和向量的size
     */
    void setup_system();
    /**
     * @brief 组装刚度矩阵和右端项，处理边界条件
     */
    void assemble_system();
    /**
     * @brief 求解，共轭梯度法
     */
    void solve();
    /**
     * @brief 输出结果
     */
    void output_results() const;
    void outputData();

    Quadrilateral_1_Element _element;/**< 模板单元类型*/
    Rectangle2D _domain;/**< 问题的定义域*/
    StructedMesh<2U> _mesh;/**< 划分的网格*/
    Eigen::SparseMatrix<double> system_matrix;/**< 刚度矩阵 */
    Eigen::VectorXd solution;/**< 结果 */
    Eigen::VectorXd system_rhs;/**< 右端项*/
};

double RightHandSide(Dofs<2U>& dofs) {
    //   double return_value = 0.0;
//   for (unsigned int i = 0; i < dim; ++i)
//       return_value += 4.0 * std::pow(p(i), 4.0);
//   return return_value;
    return 1;
}



double BoundaryFunction(Dofs<2U>& dofs)
{
    //return dofs[0]*dofs[0] + dofs[1]*dofs[1];
    return 0;
}
/**
     * @brief 返回map，存放边界上自由度的全局标号，和边界函数(仅Dirichlet)在对应自由度上的函数值。
     *
     * @param _boundary_function 边界函数
     * @param _boundary_dofs 边界上的自由度
     */
template <unsigned int DIM>
std::map<int, double> interpolateBoundaryValues(std::function<double (Dofs<DIM>&)>& _boundary_function, vector<Dofs<DIM>>& _boundary_dofs)
{
    std::map<int,double> boundary_values;
    for(Dofs<DIM>& dofs : _boundary_dofs)
    {
	int index = dofs.getGlobalIndex();
	double value = _boundary_function(dofs);
	boundary_values.insert({index, value});
    }
    return boundary_values;
}
/**
     * @brief 对刚度矩阵和右端项应用Dirichlet条件
     * 
     * @param system_matrix 刚度矩阵
     * @param system_rhs 右端项
     */
    
void applyBoundaryValues(std::map<int, double>& boundary_values, Eigen::SparseMatrix<double>& system_matrix, Eigen::VectorXd& system_rhs)
{
    for(auto& it : boundary_values)
    {
        //边界自由度的全局编号。
	int k = it.first; 
        //对应自由的边界函数值。
	double bnd_value = it.second;
	system_rhs[k] = bnd_value * system_matrix.coeffRef(k,k);
	for(Eigen::SparseMatrix<double>::InnerIterator it(system_matrix, k); it; ++it)
	{
	    int row = it.row();
	    if(row == k)
		continue;
	    system_rhs[row] -= system_matrix.coeffRef(k, row) * bnd_value;
	    system_matrix.coeffRef(k, row) = 0.0;
	    system_matrix.coeffRef(row, k) = 0.0;
	}
    }
}
/**
 * This function does the first part, creating the mesh. We create the mesh
 * which is the square [−1,1]×[−1,1].
 */
void FEM::make_grid(const vector<int>& _arg)
{
    vector<Point<2>> p = {{-2.0, -2.0}, {1.0, -2.0}, {1.0, 1.0}, {-2.0, 1.0}};
    _domain.setVertexList(p);
    StructedMesh<2> mesh(_domain, _arg);
    _mesh = mesh;
}



/**
 * Next we enumerate all the degrees of freedom and set up matrix and vector objects to hold the system data.
 *
 */
void FEM::setup_system()
{
    /*Enumerating is done by using DistributeDofs(),2-order element*/
    _mesh.distributeDofs(1);
    std::cout << "Number of degrees of freedom: " << _mesh.getTotalNumDofs() << std::endl;

    /**
     * set the sizes of system_matrix, right hand side vector ans the solution vector;
     */
    system_matrix = Eigen::SparseMatrix<double>(int(_mesh.getTotalNumDofs()),int(_mesh.getTotalNumDofs()));
    system_matrix.setZero();
    system_rhs = Eigen::VectorXd::Zero(int(_mesh.getTotalNumDofs()));
    solution = Eigen::VectorXd::Zero(int(_mesh.getTotalNumDofs()));
    std::cout << "setup system is done" << std::endl;
}


/**
 * The next step is to compute the entries of the matrix
 * and right hand side that form the linear system from which we compute the solution.
 */
void FEM::assemble_system()
{
    clock_t start,end;
    start = clock();

    /**
     * Get the right_hand_side function and boundary function
     */
    std::function<double(Dofs<2U>&)> right_hand_side = RightHandSide;
    std::function<double(Dofs<2U>&)> boundary_function = BoundaryFunction;
   


    /**
     * for each Grid, we construct a cell_matirx and cell_rhs.
     */
    const unsigned int dofs_per_cell = _element.n_Dofs();
    Eigen::MatrixXd cell_matrix(dofs_per_cell,dofs_per_cell);
    Eigen::VectorXd cell_rhs(dofs_per_cell);
    Quadrilateral_Element_GaussianInfo gauss;//加入gauss点
    gauss.LoadGaussianInfo(1);
    std::cout << "gauss points is added " << std::endl;

    /*assemble the global matrix and vector cell-by-cell.*/
    for (auto &cell : _mesh.getGridVector())
    {
    	// bind the TemplateElement with current cell.
	_element.reinit(cell);

	cell_matrix = Eigen::MatrixXd::Zero(dofs_per_cell,dofs_per_cell);
	cell_rhs    = Eigen::VectorXd::Zero(dofs_per_cell);
	for (int q_index = 0; q_index < gauss.GaussianPointNumInfo();++q_index)
        {
	    //std::cout << "GaussPnt number is: " << _element.n_GaussPnt() << std::endl;
	    Point<2> pnt = gauss.GaussianPointInfo(q_index);
	    //std::cout << "point: " << pnt << std::endl;
	    //自下面开始注意标号...一致性
	    for (int i = 0; i < dofs_per_cell; ++i)
	    {
		for (int j = 0; j < dofs_per_cell; ++j)
		{
		    double InnerProduct = 0;
		    for(int k = 0;k < _element.gradient(pnt[0],pnt[1],i+1).size();k++)
		    {
		        
			InnerProduct += _element.gradient(pnt[0],pnt[1],i+1)[k] *
			    _element.gradient(pnt[0],pnt[1],j+1)[k];
		    }
		    cell_matrix(i, j) +=(InnerProduct *
					 _element.det_Jacobi(pnt[0], pnt[1]) *
					 gauss.GaussianWeightInfo(q_index));
		}
	    }
	    double xi = pnt[0],eta = pnt[1];
	    double x = _element.Global_x(xi, eta);
	    double y = _element.Global_y(xi, eta);
	    for (int i = 0; i < dofs_per_cell; ++i)
	    {
		double x = _element.Global_x(xi,eta);
		double y = _element.Global_y(xi,eta);
        
		//vector<double> p({x,y});
		Point<2> p({x,y});

		Dofs<2> globalDof(p);
		cell_rhs(i) += ( _element.phi(xi, eta, i + 1) *
				 right_hand_side(globalDof) *
				 _element.det_Jacobi(xi, eta) *
				 gauss.GaussianWeightInfo(q_index));
	    }

	}

	for (int i = 0; i < dofs_per_cell; ++i)
	{
	    for (int j = 0; j < dofs_per_cell; ++j)
	    {
		int row = _element.getGlobalIndex(i+1);
		int column = _element.getGlobalIndex(j+1);
		system_matrix.coeffRef(row, column)+=cell_matrix(i, j);
	    }
	}
	for (int i = 0; i < dofs_per_cell; ++i)
	    system_rhs[_element.getGlobalIndex(i+1)] += cell_rhs(i);
    }
    system_matrix.makeCompressed();
 
    std::vector<Dofs<2>> boundary_dofs = _mesh.getBoundaryDofs();
    std::map<int, double> boundary_values;

    boundary_values = interpolateBoundaryValues(boundary_function, boundary_dofs);

    applyBoundaryValues(boundary_values,system_matrix, system_rhs);
    std::cout << "assemble system is done" << std::endl;
    end=clock();		//程序结束用时
    double endtime=(double)(end-start)/CLOCKS_PER_SEC;
    cout<<"assemble system takes  "<<endtime*1000<<" ms"<<endl;	//ms为单位
}



void FEM::solve()
{
    clock_t start,end;
    start = clock();

    /**
     *  stop if the norm of the residual is below τ=1e−6*||b|| where b is the right hand side vector
     */
    Eigen::ConjugateGradient<Eigen::SparseMatrix<double>,Eigen::Lower|Eigen::Upper> cg;
    cg.compute(system_matrix);
    cg.setTolerance(1e-6);
    solution = cg.solve(system_rhs);
    std::cout << "#iterations:     " << cg.iterations() << std::endl;
    std::cout << "estimated error: " << cg.error()      << std::endl;
    //std::cout << solution.size() << std::endl;
    end=clock();		//程序结束用时
    double endtime=(double)(end-start)/CLOCKS_PER_SEC;
    cout<<"solve system takes  "<<endtime*1000<<" ms"<<endl;	//ms为单位
}


void FEM::outputData()//这是要保存的名字和地址, DIM=2
{
    FILE * pFile;
    pFile = fopen("test.txt", "w");//这个用“w”是写文件，覆盖原内容，若不想覆盖则用“a”
    if (NULL == pFile)
    	printf("error");//文件打开错误
    fprintf(pFile,"%f\n%f\n%f\n%f\n", _domain.getVertex(0)[0], _domain.getVertex(1)[0], _domain.getVertex(0)[1], _domain.getVertex(3)[1]);
    fprintf(pFile,"%d\n%d\n", _mesh.getSegment()[0], _mesh.getSegment()[1]);
    for (int i = 0; i < solution.size(); i++)
    	fprintf(pFile,"%.10f\n", solution[i]);//循环输出数据并写入
    fclose(pFile);//最后一定记得关闭文件
    std::cout << "Data has been output" << std::endl; 
}

void FEM::output_results() const
{
    
}



void FEM::run(vector<int> _seg)
{
    make_grid(_seg);
    setup_system();
    assemble_system();
    solve();
    //output_results();
    outputData();
}



int main(int argc,char* argv[])
{
    vector<int> seg = {std::atoi(argv[1]), std::atoi(argv[2])};
    FEM laplace_problem;
    laplace_problem.run(seg);
    return 0;
}
